# -*- coding: utf-8 -*-
from PIL import Image
import os
import os.path as osp
import pickle

BASEDIR = osp.dirname(osp.abspath(__file__))


def load(filename):
    with open(filename, 'rb') as fo:
        data = pickle.load(fo, encoding='latin1')
    return data


for num in range(1, 1+5):
    path_file = osp.join(BASEDIR, 'cifar-data/cifar-10-python/cifar-10-batches-py', f'data_batch_{num}')
    data = load(path_file)

    out_dir = osp.join(BASEDIR, 'cifar-images', osp.basename(path_file) + '_dir')
    if not osp.exists(out_dir):
        os.makedirs(out_dir)

    print(data.keys())

    # print(data['batch_label'])
    # print(len(data['labels']))
    # print(data['data'].dtype)  # [10000 x 3 x 32 x 32] uint8
    # print(len(data['filenames']))

    num_samples = len(data['labels'])
    print(num_samples)

    for i in range(num_samples):
        label_i = data['labels'][i]
        imgArr_i = data['data'][i]
        imgArr_i = imgArr_i.reshape((3, 32, 32)).transpose((1, 2, 0))
        # print(imgArr_i.dtype)
        filename_i = data['filenames'][i]
        img_i = Image.fromarray(imgArr_i)
        out_path = osp.join(out_dir, str(label_i) + '_' + filename_i)
        img_i.save(out_path)
        print(filename_i, 'extracted.')
